# -*- coding:utf-8 -*-
import cv2
import base64
from paddleocr import PPStructure

def drawImage(img,regionList):
    for r in regionList:
        x,y,w,h = r['bbox']
        type = r['type']
        text = type
        cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
    imgStr = cv2.imencode('.jpg', img)[1].tostring()
    imgBase64 = base64.b64encode(imgStr).decode()
    return imgBase64

def getNodeInfo(img):
    regionList = []
    table_engine = PPStructure(table=False,ocr=False,layout=True,show_log=False,det_model_dir='/data/zwxu/layoutparse/paddle/det/ch/ch_PP-OCRv3_det_infer',rec_model_dir='/data/zwxu/layoutparse/paddle/rec/ch/ch_PP-OCRv3_rec_infer',table_model_dir='/data/zwxu/layoutparse/paddle/table/ch_ppstructure_mobile_v2.0_SLANet_infer',layout_model_dir='/data/zwxu/layoutparse/paddle/layout/picodet_lcnet_x1_0_fgd_layout_cdla_infer')
    result = table_engine(img)
    for line in result:
      line.pop('img')
      #左上角x，左上角y，右下角x，右下角y
      x = line['bbox'][0]
      y = line['bbox'][1]
      w = line['bbox'][2] - x
      h = line['bbox'][3] - y
      line['bbox'] = x,y,w,h 
      regionList.append(line)

    height,width,_ = img.shape
    imgBase64 = drawImage(img,regionList)
    nodeInfo = {'height':height,'width':width,'node':regionList,'base64':imgBase64}
    return nodeInfo
   
#image_path = './0722.jpg'  # 待预测图片路径
#img = cv2.imread(image_path)
#nodeInfo = getNodeInfo(img)
#print(nodeInfo)
